#!/usr/bin/env python
 
# sidekick mini web server http://sidekick.windforwings.com/2013/03/minimal-websocket-broadcast-server-in.html
 
import socket, hashlib, base64, threading
import time

class PyWSock():
    MAGIC = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
    HSHAKE_RESP = "HTTP/1.1 101 Switching Protocols\r\n" + \
                "Upgrade: websocket\r\n" + \
                "Connection: Upgrade\r\n" + \
                "Sec-WebSocket-Accept: %s\r\n" + \
                "\r\n"
    LOCK = threading.Lock()
    
    clients = []
    
    run_server = True
    
    data_to_send = None;
    data_to_send_copy = None;

    def recv_data (self, client):
        # as a simple server, we expect to receive:
        #    - all data at one go and one frame
        #    - one frame at a time
        #    - text protocol
        #    - no ping pong messages
        data = bytearray(client.recv(512))
        if(len(data) < 6):
            return b''
            #raise Exception("Error reading data")
        # FIN bit must be set to indicate end of frame
        assert(0x1 == (0xFF & data[0]) >> 7)
        # data must be a text frame
        # 0x8 (close connection) is handled with assertion failure
        assert(0x1 == (0xF & data[0]))
        
        # assert that data is masked
        assert(0x1 == (0xFF & data[1]) >> 7)
        datalen = (0x7F & data[1])
        
        #print("received data len %d" %(datalen,))
        
        str_data = ''
        if(datalen > 0):
            mask_key = data[2:6]
            masked_data = data[6:(6+datalen)]
            unmasked_data = [masked_data[i] ^ mask_key[i%4] for i in range(len(masked_data))]
            str_data = str(bytearray(unmasked_data))
        
        print("received [%s]" % (str_data,))
        return str_data

    def broadcast_resp(self, data):
        
        if not data or len(data) == 0:
            return
                           
        self.LOCK.acquire()
        for client in self.clients:
            try:
                
                data_length = len(data)
                
                # 1st byte: fin bit set. text frame bits set.
                # 2nd byte: no mask. length set in 1 byte.
                resp = bytearray([0b10000001])

                if data_length <= 125:
                    resp.append(data_length)

                elif data_length >= 126 and data_length <= 65535:
                    resp.append(  126)
                    resp.append( ( data_length >> 8 ) & 255)
                    resp.append( ( data_length      ) & 255)

                else:
                    resp.append( 127)
                    resp.append( ( data_length >> 56 ) & 255)
                    resp.append( ( data_length >> 48 ) & 255)
                    resp.append( ( data_length >> 40 ) & 255)
                    resp.append( ( data_length >> 32 ) & 255)
                    resp.append( ( data_length >> 24 ) & 255)
                    resp.append( ( data_length >> 16 ) & 255)
                    resp.append( ( data_length >>  8 ) & 255)
                    resp.append( ( data_length       ) & 255)

                # append the data bytes
                local_data = bytes(data, encoding='utf-8') 
                for d in bytearray(local_data):
                    resp.append(d)
            
                client.sendall(resp)
                # print("sending: ", resp)                    
                        
            except Exception as e:
                client.close()
                if(client in self.clients):
                    self.clients.remove(client)
                    print("error sending to a client: ", e)
                    
        self.LOCK.release()

    def split_len(self, seq, length):
        return [seq[i:i+length] for i in range(0, len(seq), length)]
 
    def parse_headers (self, data):
        headers = {}
        if data and len(data) > 0:
            lines = data.splitlines()
            for l in lines:
                l = str(l)
                parts = l.split(": ", 1)
                if len(parts) == 2:
                    headers[parts[0]] = parts[1]
            headers['code'] = lines[len(lines) - 1]
        return headers
 
    def handshake (self, client):
        print('Handshaking...')
        data = client.recv(2048)
        headers = self.parse_headers(data)
        
        if not headers or len(headers) == 0:
            return
        
        print('Got headers:')
        for k, v in headers.items():
            print (k, ':', v)
            
        key = headers["b'Sec-WebSocket-Key"]
        key = key.replace("'", '')
        encoded_key = (key+self.MAGIC).encode('utf-8')
        resp_data = self.HSHAKE_RESP % ((base64.b64encode(hashlib.sha1( encoded_key ).digest()),))
        resp_data = resp_data.replace("b'", '')
        resp_data = resp_data.replace("'", '')
        print('Response: [%s]' % (resp_data,))
        resp_data = bytes(resp_data, 'UTF-8')
        return client.send(resp_data)
     
    def handle_client (self, client, addr):
        self.handshake(client)
        try:
            while (self.run_server):            
                #data = self.recv_data(client)
                
                if self.data_to_send != self.data_to_send_copy:
                    self.broadcast_resp(self.data_to_send)
                    self.data_to_send_copy = self.data_to_send
   
                time.sleep(.1)
        except Exception as e:
            print("Exception %s" % (str(e)))
        print('Client closed: ' + str(addr))
        self.LOCK.acquire()
        if(client in self.clients):
                    self.clients.remove(client)
        self.LOCK.release()
        client.close()
        
    def start_server (self, port):
        s = socket.socket()
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        s.bind(('', port))
        s.listen(5)
        while (self.run_server):
            print ('Waiting for connection...')
            conn, addr = s.accept()
            print ('Connection from: ' + str(addr))
            threading.Thread(target = self.handle_client, args = (conn, addr)).start()
            self.LOCK.acquire()
            self.clients.append(conn)
            self.LOCK.release()
           # break;
            
    def close (self):
        print ("closing")
        self.run_server = None

    def __init__(self):
        socketThread = threading.Thread(target=self.start_server, args=(4545,))
        socketThread.start()
#        try:
#            self.start_server(4545)
#        except KeyboardInterrupt:
#            self.close()

